import json


def load_exist_movie2id(d2f, pth):
    exist_movie2id = {}
    with open(pth, "r", encoding="utf-8") as f:
        for line in f:
            mid0, mid1 = line.strip().split()
            exist_movie2id[d2f[mid0]] = int(mid1)
    return exist_movie2id


def generate_movie2id(d2f: dict, ent_info: dict, movie_id_map_pth: str):
    movie2id = load_exist_movie2id(d2f, movie_id_map_pth)
    cnt = len(movie2id)
    for k in ent_info:
        if k not in movie2id:  # 似乎有d2f中的电影被删了，导致movie2id中的值不全为ent_info.keys的子集
            movie2id[k] = cnt
            cnt += 1
    return movie2id


def generate_relation2id(rel_info):
    cnt = 0
    relation2id = {}
    for k in rel_info:
        relation2id[k] = cnt
        cnt += 1
    return relation2id


def rebuild_graph(graph: dict, movie2id: dict, relation2id: dict):
    new_graph = {}
    for head, edge in graph.items():
        h_id = movie2id[head]
        new_graph[h_id] = []
        for rel, tail in edge:
            r_id, t_id = relation2id[rel], movie2id[tail]
            new_graph[h_id].append((r_id, t_id))
    return new_graph


def write_graph(graph: dict, pth: str = "./data/Douban/kg_final.txt"):
    with open(pth, "w", encoding="utf-8") as f:
        for h_id, edge in graph.items():
            for r_id, t_id in edge:
                triplet = str(h_id) + " " + str(r_id) + " " + str(t_id) + "\n"
                f.write(triplet)


def dict_load(pth):
    with open(pth, "r", encoding="utf-8") as f:
        d = json.load(f)
    return d


def dict_save(d: dict, pth):
    d_js = json.dumps(d, indent=4)
    with open(pth, "w", encoding="utf-8") as f:
        f.write(d_js)


if __name__ == "__main__":
    d2f_pth = "./res/douban2fb.json"
    graph_pth = "./res/graph_white.json"
    info_pth = "./res/st_info_white.json"
    new_graph_pth = "./data/Douban/kg_final.txt"
    movie_id_map_pth = "./data/movie_id_map.txt"
    movie2id_pth = "./res/movie2id.json"
    relation2id_pth = "./res/relation2id.json"

    d2f = dict_load(d2f_pth)
    g = dict_load(graph_pth)
    st_info = dict_load(info_pth)

    movie2id = generate_movie2id(d2f, st_info["entity"], movie_id_map_pth)
    relation2id = generate_relation2id(st_info["relation"])
    new_g = rebuild_graph(g, movie2id, relation2id)

    dict_save(movie2id, movie2id_pth)
    dict_save(relation2id, relation2id_pth)
    write_graph(new_g, new_graph_pth)
